// Licensed to the Apache Software Foundation (ASF) under one or more
// contributor license agreements.  See the NOTICE file distributed with
// this work for additional information regarding copyright ownership.
// The ASF licenses this file to You under the Apache License, Version 2.0
// (the "License"); you may not use this file except in compliance with
// the License.  You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

pub mod channel;
pub mod reporter;
pub mod storage;

use crate::{
    channel::TxReporter,
    reporter::{ReporterConfiguration, run_reporter},
    storage::{DiskStorage, StorageManager},
};
use skywalking::{
    management::{instance::Properties, manager::Manager},
    reporter::{CollectItem, CollectItemConsume},
};
use std::{
    error::Error,
    fs, io,
    os::unix::prelude::OsStrExt,
    path::{Path, PathBuf},
    sync::Arc,
    time::Duration,
};
use tokio::{
    net::UnixListener,
    runtime::{self, Runtime},
    select,
    signal::unix::{SignalKind, signal},
    sync::{mpsc::{self, error::TrySendError}, Mutex},
    task::JoinHandle,
};
use tonic::async_trait;
use tracing::{debug, error, info, warn};

pub struct WorkerConfiguration {
    pub socket_file_path: PathBuf,
    pub heart_beat: Option<HeartBeatConfiguration>,
    pub reporter_config: ReporterConfiguration,
    pub storage_path: Option<PathBuf>,
    pub buffer_size: Option<usize>,
    pub storage_quota_mb: Option<u64>,
}

pub struct HeartBeatConfiguration {
    pub service_instance: String,
    pub service_name: String,
    pub heartbeat_period: i64,
    pub properties_report_period_factor: i64,
}

pub fn new_tokio_runtime(worker_threads: usize) -> Runtime {
    runtime::Builder::new_multi_thread()
        .thread_name("sw: worker")
        .enable_all()
        .worker_threads(worker_threads)
        .build()
        .unwrap()
}

pub async fn start_worker(config: WorkerConfiguration) -> anyhow::Result<()> {
    debug!("Starting worker...");

    let socket_file = config.socket_file_path;

    let _guard = WorkerExitGuard(socket_file.clone());

    let mut sig_term = signal(SignalKind::terminate())?;
    let mut sig_int = signal(SignalKind::interrupt())?;

    let buffer_size = config.buffer_size.unwrap_or(1000);
    
    let storage_manager = if let Some(storage_path) = &config.storage_path {
        info!(?storage_path, "Initializing disk storage");
        if !storage_path.exists() {
            fs::create_dir_all(storage_path)?;
        }
        
        let quota = config.storage_quota_mb.unwrap_or_else(|| {
            let available_space = get_available_disk_space(storage_path);
            let ten_percent = available_space.saturating_div(10);
            let one_gb = 1024;
            std::cmp::min(ten_percent, one_gb)
        });
        
        info!(quota_mb = quota, "Setting storage quota");
        
        let storage = DiskStorage::new(storage_path.clone(), quota * 1024 * 1024)?;
        Some(Arc::new(Mutex::new(StorageManager::new(storage))))
    } else {
        debug!("Disk storage not configured, running in memory-only mode");
        None
    };

    let fut = async move {
        debug!(?socket_file, "Bind unix stream");
        let listener = UnixListener::bind(&socket_file)?;
        change_permission(socket_file, 0o777);

        let (tx, rx) = mpsc::channel::<CollectItem>(buffer_size);
        let _tx_reporter = tx.clone();
        
        let _health_check_task = spawn_health_check_task(storage_manager.clone());
        
        // 创建发送者副本用于心跳服务和存储恢复任务
        let tx_for_listener = tx.clone();
        let tx_heartbeat = tx.clone();
        let tx_recovery = tx.clone();
        let storage_for_listener = storage_manager.clone();
        
        tokio::spawn(async move {
            loop {
                match listener.accept().await {
                    Ok((mut stream, _addr)) => {
                        let tx = tx_for_listener.clone();
                        let storage_mgr = storage_for_listener.clone();

                        tokio::spawn(async move {
                            debug!("Entering channel_receive loop");

                            loop {
                                let r = match channel::channel_receive(&mut stream).await {
                                    Err(err) => match err.downcast_ref::<io::Error>() {
                                        Some(e) if e.kind() == io::ErrorKind::UnexpectedEof => {
                                            debug!("Leaving channel_receive loop");
                                            return;
                                        }
                                        _ => {
                                            error!(?err, "channel_receive failed");
                                            continue;
                                        }
                                    },
                                    Ok(i) => i,
                                };

                                if let Err(err) = tx.try_send(r.clone()) {
                                    if matches!(err, TrySendError::Full(_)) {
                                        warn!("Channel buffer full, attempting to store data on disk");
                                        if let Some(storage_mgr) = &storage_mgr {
                                            if let Err(store_err) = storage_mgr.lock().await.store(r).await {
                                                error!(?store_err, "Failed to store data on disk");
                                            } else {
                                                debug!("Data successfully stored on disk");
                                            }
                                        } else {
                                            error!("Channel buffer full and disk storage not available, data lost");
                                        }
                                    } else {
                                        error!(?err, "Send collect item failed");
                                        return;
                                    }
                                }
                            }
                        });
                    }
                    Err(err) => {
                        error!(?err, "Accept failed");
                    }
                }
            }
        });

        if let Some(heart_beat_config) = config.heart_beat {
            report_properties_and_keep_alive(heart_beat_config, TxReporter(tx_heartbeat));
        }

        let recovery_task = if let Some(storage_mgr) = storage_manager.clone() {
            spawn_storage_recovery_task(storage_mgr, tx_recovery)
        } else {
            tokio::spawn(async { Ok(()) })
        };

        // 创建Consumer，并使用它来创建ConsumerWithRecovery
        let consumer = Consumer(rx);
        let consumer_with_recovery = ConsumerWithRecovery::new(consumer, recovery_task);

        let reporter_result = run_reporter(config.reporter_config, (), consumer_with_recovery).await;
        
        reporter_result
    };

    select! {
        _ = sig_term.recv() => {
            info!("Received SIGTERM, gracefully shutting down...");
            graceful_shutdown().await;
        }
        _ = sig_int.recv() => {
            info!("Received SIGINT, gracefully shutting down...");
            graceful_shutdown().await;
        }
        r = fut => {
            if let Err(e) = r {
                error!(?e, "Worker failed");
            }
        }
    }

    info!("Skywalking worker shutdown complete");

    Ok(())
}

async fn graceful_shutdown() {
    info!("Starting graceful shutdown");
    let shutdown_timeout = Duration::from_secs(10);
    
    tokio::select! {
        _ = tokio::time::sleep(shutdown_timeout) => {
            warn!("Shutdown timeout reached, forcing exit");
        }
    }
}

fn get_available_disk_space(path: &Path) -> u64 {
    match fs_extra::dir::get_size(path) {
        Ok(dir_size) => {
            // 使用std::fs获取总空间而不是fs_extra
            let total_space = match std::fs::metadata(path) {
                Ok(_) => {
                    #[cfg(target_os = "linux")]
                    {
                        match std::fs::metadata(path)
                            .and_then(|_| std::io::Result::Ok(nix::sys::statvfs::statvfs(path)))
                        {
                            Ok(Ok(stat)) => {
                                let block_size = stat.block_size() as u64;
                                let blocks = stat.blocks() as u64;
                                block_size * blocks
                            }
                            _ => 1024 * 1024 * 1024, // 默认1GB
                        }
                    }
                    
                    #[cfg(not(target_os = "linux"))]
                    {
                        1024 * 1024 * 1024 // 默认1GB，非Linux平台
                    }
                }
                Err(_) => 1024 * 1024 * 1024,  // 默认1GB
            };
            
            let used_space = dir_size;
            let available = if total_space > used_space {
                total_space - used_space
            } else {
                1024 * 1024  // 默认1MB
            };
            available / (1024 * 1024) // 转换为MB
        }
        Err(e) => {
            warn!(?e, ?path, "Failed to get disk space, assuming 1GB");
            1024 // 默认1GB
        }
    }
}

fn spawn_health_check_task(storage_manager: Option<Arc<Mutex<StorageManager>>>) -> JoinHandle<()> {
    let storage = storage_manager.clone();
    
    tokio::spawn(async move {
        let check_interval = Duration::from_secs(60);
        
        loop {
            tokio::time::sleep(check_interval).await;
            
            if let Some(storage) = &storage {
                match storage.lock().await.check_health().await {
                    Ok(stats) => {
                        info!(
                            files = stats.file_count,
                            total_size_mb = stats.total_size / (1024 * 1024),
                            usage_percent = stats.usage_percent,
                            "Storage health check"
                        );
                    }
                    Err(e) => {
                        error!(?e, "Storage health check failed");
                    }
                }
            }
            
            let memory_usage = get_memory_usage();
            info!(
                mem_used_mb = memory_usage.used_mb,
                mem_total_mb = memory_usage.total_mb,
                mem_percent = memory_usage.percent,
                "Memory usage"
            );
        }
    })
}

fn spawn_storage_recovery_task(
    storage_manager: Arc<Mutex<StorageManager>>,
    tx: mpsc::Sender<CollectItem>
) -> JoinHandle<anyhow::Result<()>> {
    tokio::spawn(async move {
        let check_interval = Duration::from_secs(5);
        let mut consecutive_successes = 0;
        
        loop {
            tokio::time::sleep(check_interval).await;
            
            if tx.capacity() > 0 {
                let batch_size = tx.capacity().min(50);
                
                match storage_manager.lock().await.recover_batch(batch_size).await {
                    Ok(items) => {
                        if !items.is_empty() {
                            let recovered = items.len();
                            consecutive_successes += 1;
                            
                            // 先计算剩余项目数，然后再记录日志，确保跨await点的格式化是安全的
                            let remaining = storage_manager.lock().await.stored_items_count().await;
                            info!(recovered, remaining, "Recovered data from storage");
                            
                            for item in items {
                                if let Err(e) = tx.send(item).await {
                                    error!("Failed to send recovered item: {:?}", e);
                                    break;
                                }
                            }
                            
                            let items_remaining = storage_manager.lock().await.stored_items_count().await;
                            if consecutive_successes >= 5 && items_remaining == 0 {
                                debug!("All data recovered, pausing recovery task");
                                tokio::time::sleep(Duration::from_secs(30)).await;
                                consecutive_successes = 0;
                            }
                        } else {
                            tokio::time::sleep(Duration::from_secs(30)).await;
                        }
                    }
                    Err(e) => {
                        error!("Failed to recover data from storage: {:?}", e);
                        consecutive_successes = 0;
                        tokio::time::sleep(Duration::from_secs(60)).await;
                    }
                }
            } else {
                debug!("Send channel full, waiting before recovering more data");
                consecutive_successes = 0;
                tokio::time::sleep(Duration::from_secs(5)).await;
            }
        }
    })
}

struct ConsumerWithRecovery {
    consumer: Consumer,
    #[allow(dead_code)]
    recovery_task: JoinHandle<anyhow::Result<()>>,
}

impl ConsumerWithRecovery {
    fn new(consumer: Consumer, recovery_task: JoinHandle<anyhow::Result<()>>) -> Self {
        Self { consumer, recovery_task }
    }
}

// 实现CollectItemConsume trait
#[async_trait]
impl CollectItemConsume for ConsumerWithRecovery {
    async fn consume(&mut self) -> Result<Option<CollectItem>, Box<dyn Error + Send>> {
        self.consumer.consume().await
    }
    
    async fn try_consume(&mut self) -> Result<Option<CollectItem>, Box<dyn Error + Send>> {
        self.consumer.try_consume().await
    }
}

// 移除对mpsc::Receiver<CollectItem>的实现
// 保留Consumer结构体
struct Consumer(mpsc::Receiver<CollectItem>);

#[async_trait]
impl CollectItemConsume for Consumer {
    async fn consume(&mut self) -> Result<Option<CollectItem>, Box<dyn Error + Send>> {
        Ok(self.0.recv().await)
    }

    async fn try_consume(&mut self) -> Result<Option<CollectItem>, Box<dyn Error + Send>> {
        Ok(self.0.try_recv().ok())
    }
}

struct MemoryUsage {
    used_mb: u64,
    total_mb: u64,
    percent: f64,
}

#[cfg(target_os = "linux")]
fn get_memory_usage() -> MemoryUsage {
    match procfs::Meminfo::new() {
        Ok(meminfo) => {
            let total_kb = meminfo.mem_total;
            let free_kb = meminfo.mem_free;
            
            // procfs 0.15.1版本中，buffers和cached实际是u64而不是Option<u64>
            let buffers_kb = meminfo.buffers;
            let cached_kb = meminfo.cached;
            
            let used_kb = total_kb - free_kb - buffers_kb - cached_kb;
            let total_mb = total_kb / 1024;
            let used_mb = used_kb / 1024;
            let percent = if total_kb > 0 {
                (used_kb as f64 / total_kb as f64) * 100.0
            } else {
                0.0
            };
            
            MemoryUsage {
                used_mb,
                total_mb,
                percent,
            }
        }
        Err(_) => default_memory_usage(),
    }
}

#[cfg(not(target_os = "linux"))]
fn get_memory_usage() -> MemoryUsage {
    default_memory_usage()
}

fn default_memory_usage() -> MemoryUsage {
    MemoryUsage {
        used_mb: 0,
        total_mb: 0,
        percent: 0.0,
    }
}

struct WorkerExitGuard(PathBuf);

impl Drop for WorkerExitGuard {
    fn drop(&mut self) {
        let &mut Self(ref socket_file) = self;
        info!(?socket_file, "Remove socket file");
        if let Err(err) = fs::remove_file(socket_file) {
            error!(?err, "Remove socket file failed");
        }
    }
}

fn report_properties_and_keep_alive(config: HeartBeatConfiguration, reporter: TxReporter) {
    let manager = Manager::new(&*config.service_name, &*config.service_instance, reporter);

    manager.report_and_keep_alive(
        || {
            let mut props = Properties::new();
            props.insert_os_info();
            props.update(Properties::KEY_LANGUAGE, "php");
            props.update(Properties::KEY_PROCESS_NO, unsafe {
                libc::getppid().to_string()
            });
            debug!(?props, "Report instance properties");
            props
        },
        Duration::from_secs(config.heartbeat_period as u64),
        config.properties_report_period_factor as usize,
    );
}

fn change_permission(f: impl AsRef<Path>, mode: libc::mode_t) {
    let f = f.as_ref().as_os_str().as_bytes();
    let mut path = Vec::with_capacity(f.len() + 1);
    path.extend_from_slice(f);
    path.push(b'\0');
    unsafe {
        libc::chmod(path.as_ptr().cast(), mode);
    }
}
